"""
Bias Predictor Functions for All Observation Types

This module implements comprehensive bias predictor functions used in atmospheric 
data assimilation for computing systematic bias corrections. Predictors capture
environmental, geometric, and instrumental factors that contribute to observation biases.

Supports all GSI observation types:
- Satellite radiances (IR, MW, hyperspectral)
- Conventional observations (radiosondes, aircraft, surface)
- GPS radio occultation and precipitable water
- Scatterometer winds and altimeter data
- Lightning detection and precipitation radar
"""
module BiasPredictors

using LinearAlgebra
using Statistics
using Dates
using ..ObservationTypes
using ..BiasCorrector

export evaluate_all_predictors, get_predictor_names
export ConstantPredictor, LapseRatePredictor, LatitudePredictor
export EmissivityPredictor, ScanAnglePredictor, CloudPredictor
export SurfaceTemperaturePredictor, ThicknessPredictor
export AtmosphericPredictor, GeometricPredictor, InstrumentalPredictor
export RadianceBiasPredictors, ConventionalBiasPredictors
export GPSBiasPredictors, ScatterometerBiasPredictors

"""
Constant bias predictor - provides baseline offset correction
"""
struct ConstantPredictor{T} <: BiasPredictor{T}
    value::T
    
    ConstantPredictor{T}(value=one(T)) where T = new{T}(value)
end

function evaluate_predictor(pred::ConstantPredictor{T}, obs::Observation{T}) where T
    return pred.value
end

"""
Atmospheric lapse rate predictor for temperature-dependent biases
"""
struct LapseRatePredictor{T} <: BiasPredictor{T}
    reference_level::T  # hPa
    power::Int          # 1 for linear, 2 for quadratic
    scaling_factor::T
    
    function LapseRatePredictor{T}(ref_level=T(200), power=1, scale=one(T)) where T
        new{T}(ref_level, power, scale)
    end
end

function evaluate_predictor(pred::LapseRatePredictor{T}, obs::Observation{T}) where T
    if hasfield(typeof(obs), :atmospheric_profile)
        profile = obs.atmospheric_profile
        if hasfield(typeof(profile), :temperature) && hasfield(typeof(profile), :pressure)
            # Find temperature at reference level
            ref_temp = interpolate_profile(profile.temperature, profile.pressure, pred.reference_level)
            surface_temp = profile.temperature[1]  # Assume first level is surface
            
            lapse_rate = (surface_temp - ref_temp) / log(profile.pressure[1] / pred.reference_level)
            return pred.scaling_factor * (lapse_rate ^ pred.power)
        end
    end
    
    # Default fallback using standard atmosphere
    return pred.scaling_factor * T(6.5) ^ pred.power  # Standard lapse rate K/km
end

"""
Latitude-dependent bias predictor (cosine and sine components)
"""
struct LatitudePredictor{T} <: BiasPredictor{T}
    component::Symbol  # :cosine, :sine, :linear
    scaling_factor::T
    
    LatitudePredictor{T}(comp::Symbol, scale=one(T)) where T = new{T}(comp, scale)
end

function evaluate_predictor(pred::LatitudePredictor{T}, obs::Observation{T}) where T
    lat_rad = deg2rad(obs.location.latitude)
    
    if pred.component == :cosine
        return pred.scaling_factor * cos(lat_rad)
    elseif pred.component == :sine
        return pred.scaling_factor * sin(lat_rad)
    elseif pred.component == :linear
        return pred.scaling_factor * obs.location.latitude / T(90)
    else
        error("Unknown latitude component: $(pred.component)")
    end
end

"""
Surface emissivity predictor for microwave radiances
"""
struct EmissivityPredictor{T} <: BiasPredictor{T}
    channel_frequency::T  # GHz
    surface_type_weights::Dict{Symbol, T}
    polarization::Symbol  # :vertical, :horizontal, :mixed
    
    function EmissivityPredictor{T}(freq, weights=Dict{Symbol, T}(), pol=:vertical) where T
        new{T}(freq, weights, pol)
    end
end

function evaluate_predictor(pred::EmissivityPredictor{T}, obs::Observation{T}) where T
    if obs isa RadianceObservation
        # Surface emissivity model based on frequency and surface type
        surface_type = get_surface_type(obs.location)
        base_emissivity = get(pred.surface_type_weights, surface_type, T(0.9))
        
        # Frequency-dependent correction
        freq_factor = exp(-pred.channel_frequency / T(100))  # Simplified model
        
        # Polarization correction
        pol_factor = pred.polarization == :vertical ? T(1.02) : 
                    pred.polarization == :horizontal ? T(0.98) : one(T)
        
        return base_emissivity * freq_factor * pol_factor
    end
    
    return T(0.9)  # Default emissivity
end

"""
Satellite scan angle predictor for geometric bias correction
"""
struct ScanAnglePredictor{T} <: BiasPredictor{T}
    power::Int  # 1 for linear, 2 for quadratic
    max_angle::T  # degrees
    asymmetry_factor::T  # accounts for east/west asymmetry
    
    function ScanAnglePredictor{T}(power=2, max_angle=T(58), asymmetry=zero(T)) where T
        new{T}(power, max_angle, asymmetry)
    end
end

function evaluate_predictor(pred::ScanAnglePredictor{T}, obs::Observation{T}) where T
    if obs isa RadianceObservation && hasfield(typeof(obs), :scan_angle)
        # Normalize scan angle
        normalized_angle = obs.scan_angle / pred.max_angle
        
        # Apply power law
        angle_term = abs(normalized_angle) ^ pred.power
        
        # Add asymmetry correction
        asymmetry_term = pred.asymmetry_factor * sign(normalized_angle) * normalized_angle
        
        return angle_term + asymmetry_term
    end
    
    return zero(T)
end

"""
Cloud liquid/ice water predictor for atmospheric scattering effects
"""
struct CloudPredictor{T} <: BiasPredictor{T}
    cloud_type::Symbol  # :liquid, :ice, :total
    vertical_weighting::Vector{T}  # pressure level weights
    nonlinearity_factor::T
    
    function CloudPredictor{T}(ctype=:total, weights=T[], nonlin=T(0.5)) where T
        new{T}(ctype, weights, nonlin)
    end
end

function evaluate_predictor(pred::CloudPredictor{T}, obs::Observation{T}) where T
    if hasfield(typeof(obs), :atmospheric_profile)
        profile = obs.atmospheric_profile
        
        if hasfield(typeof(profile), :cloud_liquid_water) || 
           hasfield(typeof(profile), :cloud_ice_water)
            
            cloud_amount = if pred.cloud_type == :liquid && hasfield(typeof(profile), :cloud_liquid_water)
                sum(profile.cloud_liquid_water)
            elseif pred.cloud_type == :ice && hasfield(typeof(profile), :cloud_ice_water)
                sum(profile.cloud_ice_water)
            elseif pred.cloud_type == :total
                liquid = hasfield(typeof(profile), :cloud_liquid_water) ? 
                        sum(profile.cloud_liquid_water) : zero(T)
                ice = hasfield(typeof(profile), :cloud_ice_water) ? 
                     sum(profile.cloud_ice_water) : zero(T)
                liquid + ice
            else
                zero(T)
            end
            
            # Apply vertical weighting if provided
            if !isempty(pred.vertical_weighting)
                cloud_amount = dot(cloud_amount, pred.vertical_weighting)
            end
            
            # Nonlinear transformation
            return cloud_amount * (one(T) + pred.nonlinearity_factor * cloud_amount)
        end
    end
    
    return zero(T)
end

"""
Surface temperature predictor for land/sea temperature effects
"""
struct SurfaceTemperaturePredictor{T} <: BiasPredictor{T}
    reference_temperature::T  # K
    seasonal_amplitude::T
    diurnal_amplitude::T
    surface_type_scaling::Dict{Symbol, T}
    
    function SurfaceTemperaturePredictor{T}(ref_temp=T(288), seasonal=T(10), 
                                          diurnal=T(5), scaling=Dict{Symbol, T}()) where T
        new{T}(ref_temp, seasonal, diurnal, scaling)
    end
end

function evaluate_predictor(pred::SurfaceTemperaturePredictor{T}, obs::Observation{T}) where T
    if hasfield(typeof(obs), :surface_conditions)
        surf_temp = obs.surface_conditions.temperature
        
        # Temperature anomaly from reference
        temp_anomaly = surf_temp - pred.reference_temperature
        
        # Surface type scaling
        surface_type = get_surface_type(obs.location)
        scaling = get(pred.surface_type_scaling, surface_type, one(T))
        
        return scaling * temp_anomaly / T(10)  # Normalize to ~10K scale
    end
    
    return zero(T)
end

"""
Atmospheric thickness predictor for hydrostatic effects
"""
struct ThicknessPredictor{T} <: BiasPredictor{T}
    pressure_levels::Tuple{T, T}  # hPa (bottom, top)
    reference_thickness::T        # m
    
    function ThicknessPredictor{T}(levels=(T(1000), T(500)), ref_thick=T(5400)) where T
        new{T}(levels, ref_thick)
    end
end

function evaluate_predictor(pred::ThicknessPredictor{T}, obs::Observation{T}) where T
    if hasfield(typeof(obs), :atmospheric_profile)
        profile = obs.atmospheric_profile
        
        if hasfield(typeof(profile), :geopotential_height) && hasfield(typeof(profile), :pressure)
            # Find geopotential heights at specified pressure levels
            h_bottom = interpolate_profile(profile.geopotential_height, profile.pressure, pred.pressure_levels[1])
            h_top = interpolate_profile(profile.geopotential_height, profile.pressure, pred.pressure_levels[2])
            
            thickness = h_top - h_bottom
            thickness_anomaly = thickness - pred.reference_thickness
            
            return thickness_anomaly / T(100)  # Normalize to ~100m scale
        end
    end
    
    return zero(T)
end

"""
Comprehensive radiance bias predictors for satellite instruments
"""
struct RadianceBiasPredictors{T} <: BiasPredictor{T}
    predictors::Vector{BiasPredictor{T}}
    channel_specific::Dict{Int, Vector{BiasPredictor{T}}}
    instrument_type::Symbol  # :infrared, :microwave, :hyperspectral
    
    function RadianceBiasPredictors{T}(inst_type::Symbol) where T
        predictors = BiasPredictor{T}[]
        
        # Universal radiance predictors
        push!(predictors, ConstantPredictor{T}())
        push!(predictors, LapseRatePredictor{T}(T(200), 1))
        push!(predictors, LapseRatePredictor{T}(T(200), 2))
        push!(predictors, LatitudePredictor{T}(:cosine))
        push!(predictors, LatitudePredictor{T}(:sine))
        push!(predictors, ScanAnglePredictor{T}(1))
        push!(predictors, ScanAnglePredictor{T}(2))
        
        if inst_type == :microwave
            push!(predictors, EmissivityPredictor{T}(T(50)))
            push!(predictors, CloudPredictor{T}(:liquid))
            push!(predictors, CloudPredictor{T}(:ice))
        elseif inst_type == :infrared
            push!(predictors, CloudPredictor{T}(:total))
            push!(predictors, SurfaceTemperaturePredictor{T}())
        end
        
        push!(predictors, ThicknessPredictor{T}())
        
        new{T}(predictors, Dict{Int, Vector{BiasPredictor{T}}}(), inst_type)
    end
end

function evaluate_predictor(pred::RadianceBiasPredictors{T}, obs::RadianceObservation{T}) where T
    results = T[]
    
    # Evaluate common predictors
    for predictor in pred.predictors
        push!(results, evaluate_predictor(predictor, obs))
    end
    
    # Evaluate channel-specific predictors if available
    if haskey(pred.channel_specific, obs.channel)
        for predictor in pred.channel_specific[obs.channel]
            push!(results, evaluate_predictor(predictor, obs))
        end
    end
    
    return results
end

"""
Conventional observation bias predictors
"""
struct ConventionalBiasPredictors{T} <: BiasPredictor{T}
    predictors::Vector{BiasPredictor{T}}
    observation_type::Symbol  # :radiosonde, :aircraft, :surface, :ship, :buoy
    
    function ConventionalBiasPredictors{T}(obs_type::Symbol) where T
        predictors = BiasPredictor{T}[]
        
        # Common conventional predictors
        push!(predictors, ConstantPredictor{T}())
        push!(predictors, LatitudePredictor{T}(:linear))
        
        if obs_type in [:radiosonde, :aircraft]
            push!(predictors, LapseRatePredictor{T}(T(500), 1))
            push!(predictors, ThicknessPredictor{T}((T(850), T(200))))
        end
        
        if obs_type in [:surface, :ship, :buoy]
            push!(predictors, SurfaceTemperaturePredictor{T}())
        end
        
        new{T}(predictors, obs_type)
    end
end

function evaluate_predictor(pred::ConventionalBiasPredictors{T}, obs::ConventionalObservation{T}) where T
    return [evaluate_predictor(predictor, obs) for predictor in pred.predictors]
end

"""
GPS radio occultation bias predictors
"""
struct GPSBiasPredictors{T} <: BiasPredictor{T}
    predictors::Vector{BiasPredictor{T}}
    
    function GPSBiasPredictors{T}() where T
        predictors = BiasPredictor{T}[
            ConstantPredictor{T}(),
            LatitudePredictor{T}(:cosine),
            LatitudePredictor{T}(:sine),
            # GPS-specific predictors would go here
        ]
        
        new{T}(predictors)
    end
end

function evaluate_predictor(pred::GPSBiasPredictors{T}, obs::GPSObservation{T}) where T
    return [evaluate_predictor(predictor, obs) for predictor in pred.predictors]
end

"""
Scatterometer wind bias predictors
"""
struct ScatterometerBiasPredictors{T} <: BiasPredictor{T}
    predictors::Vector{BiasPredictor{T}}
    
    function ScatterometerBiasPredictors{T}() where T
        predictors = BiasPredictor{T}[
            ConstantPredictor{T}(),
            LatitudePredictor{T}(:cosine),
            SurfaceTemperaturePredictor{T}(),
            # Wind speed and direction predictors would go here
        ]
        
        new{T}(predictors)
    end
end

function evaluate_predictor(pred::ScatterometerBiasPredictors{T}, obs::ScatterometerObservation{T}) where T
    return [evaluate_predictor(predictor, obs) for predictor in pred.predictors]
end

"""
Evaluate all relevant predictors for a given observation
"""
function evaluate_all_predictors(obs::Observation{T}, 
                                predictors::Vector{BiasPredictor{T}}) where T
    results = T[]
    
    for predictor in predictors
        if isa(predictor, RadianceBiasPredictors) && obs isa RadianceObservation
            append!(results, evaluate_predictor(predictor, obs))
        elseif isa(predictor, ConventionalBiasPredictors) && obs isa ConventionalObservation
            append!(results, evaluate_predictor(predictor, obs))
        elseif isa(predictor, GPSBiasPredictors) && obs isa GPSObservation
            append!(results, evaluate_predictor(predictor, obs))
        elseif isa(predictor, ScatterometerBiasPredictors) && obs isa ScatterometerObservation
            append!(results, evaluate_predictor(predictor, obs))
        else
            # Generic predictor evaluation
            push!(results, evaluate_predictor(predictor, obs))
        end
    end
    
    return results
end

"""
Get predictor names for documentation and debugging
"""
function get_predictor_names(predictors::Vector{BiasPredictor{T}}) where T
    names = String[]
    
    for (i, predictor) in enumerate(predictors)
        if isa(predictor, ConstantPredictor)
            push!(names, "constant")
        elseif isa(predictor, LapseRatePredictor)
            power_str = predictor.power == 1 ? "" : "_squared"
            push!(names, "lapse_rate$(power_str)")
        elseif isa(predictor, LatitudePredictor)
            push!(names, "latitude_$(predictor.component)")
        elseif isa(predictor, EmissivityPredictor)
            push!(names, "emissivity_$(predictor.polarization)")
        elseif isa(predictor, ScanAnglePredictor)
            power_str = predictor.power == 1 ? "" : "_squared"
            push!(names, "scan_angle$(power_str)")
        elseif isa(predictor, CloudPredictor)
            push!(names, "cloud_$(predictor.cloud_type)")
        elseif isa(predictor, SurfaceTemperaturePredictor)
            push!(names, "surface_temperature")
        elseif isa(predictor, ThicknessPredictor)
            push!(names, "thickness_$(predictor.pressure_levels[1])_$(predictor.pressure_levels[2])")
        else
            push!(names, "predictor_$i")
        end
    end
    
    return names
end

# Helper functions for complex predictor evaluations
function interpolate_profile(values::Vector{T}, pressure::Vector{T}, target_pressure::T) where T
    if target_pressure >= pressure[1]
        return values[1]
    elseif target_pressure <= pressure[end]
        return values[end]
    end
    
    # Linear interpolation in log-pressure coordinates
    for i in 1:(length(pressure)-1)
        if pressure[i+1] <= target_pressure <= pressure[i]
            log_p_ratio = log(target_pressure/pressure[i+1]) / log(pressure[i]/pressure[i+1])
            return values[i+1] + log_p_ratio * (values[i] - values[i+1])
        end
    end
    
    return values[1]  # Fallback
end

function get_surface_type(location::GeographicLocation{T}) where T
    # Simplified surface type determination
    # Full implementation would use land/sea mask and surface classification
    if abs(location.latitude) > 60
        return :ice
    elseif location.latitude > 30 || location.latitude < -30
        return :land
    else
        return :ocean
    end
end

end # module BiasPredictors